import json
import numpy as np

def get_sorted_indices(scores):
    """
    返回按分数降序排列的索引列表
    例如，如果scores=[0.3, 0.8, 0.5]，则返回[1, 2, 0]
    表示排名第1的是索引1，排名第2的是索引2，排名第3的是索引0
    """
    return sorted(range(len(scores)), key=lambda i: -scores[i])

# 从eval_data.json中提取的gold_scores
q1_gold_scores = [4.5, 2.1, 3.8, 1.0, 5.0, 2.5, 3.0, 1.5, 4.0, 3.5]
q2_gold_scores = [2.0, 5.0, 1.0, 3.5, 4.0, 1.5, 3.0, 4.5, 2.5, 1.2]

# 理想排序（按gold_scores降序排列的索引）
q1_ideal_indices = get_sorted_indices(q1_gold_scores)  # [4, 0, 8, 2, 9, 6, 5, 1, 7, 3]
q2_ideal_indices = get_sorted_indices(q2_gold_scores)  # [1, 7, 4, 3, 6, 8, 0, 5, 9, 2]

# 生成qwen模型的分数（非常接近理想排序）
qwen_q1_scores = [0.0] * 10
qwen_q2_scores = [0.0] * 10

# 为qwen分配接近理想排序的分数（0.9到0.1的递减）
for i, idx in enumerate(q1_ideal_indices):
    qwen_q1_scores[idx] = 0.95 - i * 0.08

for i, idx in enumerate(q2_ideal_indices):
    qwen_q2_scores[idx] = 0.95 - i * 0.08

# 生成bge模型的分数（与理想排序有一定差异）
bge_q1_scores = [0.0] * 10
bge_q2_scores = [0.0] * 10

# 打乱理想排序，模拟bge效果不如qwen
bge_q1_ideal = q1_ideal_indices.copy()
bge_q2_ideal = q2_ideal_indices.copy()

# 故意交换一些位置，使排序与理想排序有差异
bge_q1_ideal[0], bge_q1_ideal[2] = bge_q1_ideal[2], bge_q1_ideal[0]  # 交换前两名
bge_q1_ideal[3], bge_q1_ideal[5] = bge_q1_ideal[5], bge_q1_ideal[3]  # 交换中间位置
bge_q1_ideal[7], bge_q1_ideal[8] = bge_q1_ideal[8], bge_q1_ideal[7]  # 交换靠后位置

bge_q2_ideal[0], bge_q2_ideal[1] = bge_q2_ideal[1], bge_q2_ideal[0]  # 交换前两名
bge_q2_ideal[2], bge_q2_ideal[4] = bge_q2_ideal[4], bge_q2_ideal[2]  # 交换中间位置
bge_q2_ideal[6], bge_q2_ideal[8] = bge_q2_ideal[8], bge_q2_ideal[6]  # 交换靠后位置

# 为bge分配分数
for i, idx in enumerate(bge_q1_ideal):
    bge_q1_scores[idx] = 0.92 - i * 0.08

for i, idx in enumerate(bge_q2_ideal):
    bge_q2_scores[idx] = 0.92 - i * 0.08

# 四舍五入到2位小数，使数据更整洁
bge_q1_scores = [round(s, 2) for s in bge_q1_scores]
bge_q2_scores = [round(s, 2) for s in bge_q2_scores]
qwen_q1_scores = [round(s, 2) for s in qwen_q1_scores]
qwen_q2_scores = [round(s, 2) for s in qwen_q2_scores]

# 构建数据结构
bge_data = {
    "question_1": {
        "scores": bge_q1_scores,
        "rankings": get_sorted_indices(bge_q1_scores)
    },
    "question_2": {
        "scores": bge_q2_scores,
        "rankings": get_sorted_indices(bge_q2_scores)
    }
}

qwen_data = {
    "question_1": {
        "scores": qwen_q1_scores,
        "rankings": get_sorted_indices(qwen_q1_scores)
    },
    "question_2": {
        "scores": qwen_q2_scores,
        "rankings": get_sorted_indices(qwen_q2_scores)
    }
}

# 打印理想排序和模型排序的对比
print("问题1理想排序:", q1_ideal_indices)
print("Qwen排序:", get_sorted_indices(qwen_q1_scores))
print("BGE排序:", get_sorted_indices(bge_q1_scores))
print("\n问题2理想排序:", q2_ideal_indices)
print("Qwen排序:", get_sorted_indices(qwen_q2_scores))
print("BGE排序:", get_sorted_indices(bge_q2_scores))

# 使用标准JSON格式，适当缩进
with open('data/raw/bge_scores_sample.json', 'w', encoding='utf-8') as f:
    json.dump(bge_data, f, ensure_ascii=False, indent=2)

with open('data/raw/qwen_scores_sample.json', 'w', encoding='utf-8') as f:
    json.dump(qwen_data, f, ensure_ascii=False, indent=2)

print('\n已生成更经典的样例数据！Qwen排序更接近理想排序，效果更好。') 